/// 表示字典树节点中存储值的类型别名
/// 对应 C 语言中的 `TrieValue`，使用泛型参数 T 增强类型安全性
pub type TrieValue<T> = Option<T>;


use std::cell::RefCell;
use std::rc::Rc;

/// 字典树节点结构
/// 对应 C 语言中的 `TrieNode`，使用泛型参数 T 增强类型安全性
#[derive(Debug)]
pub struct TrieNode<T> {
    /// 节点存储的值（None 表示无数据）
    pub data: TrieValue<T>,

    /// 节点引用计数器
    /// 记录该节点被多少子节点或操作引用，用于内存管理
    pub use_count: usize,

    /// 子节点指针数组
    /// 使用 Rc<RefCell<..>> 实现共享所有权和内部可变性
    /// 数组长度固定为 256，对应 ASCII 字符集
    pub next: [Option<Rc<RefCell<TrieNode<T>>>>; 256],
}


/// 字典树结构，包含根节点的指针
/// 对应C语言中的 `Trie` 结构体，使用泛型参数 T 增强类型安全性
#[derive(Debug)]
pub struct Trie<T> {
    /// 根节点指针，使用 Rc<RefCell<..>> 实现共享所有权和内部可变性
    pub root_node: Option<Rc<RefCell<TrieNode<T>>>>,
    /// 存储的条目数量
    /// 存储的条目数量
    size: usize,
}

impl<T> Trie<T> {
    /// 创建新的空字典树实例
    ///
    /// # Returns
    /// 返回已初始化的 Trie 结构体实例，root_node 字段为 None
    pub fn new() -> Self {
        Self {
            root_node: None,
            size: 0,
        }
    }
}

impl<T> Trie<T> {
    /// 查找二进制键对应的终端节点
    ///
    /// # 参数
    /// - `key`: 二进制键值，以字节切片形式传递
    ///
    /// # 返回值
    /// 返回 Option<Rc<RefCell<TrieNode<T>>>>：
    /// - Some 包含找到的终端节点智能指针
    /// - None 表示键不存在于字典树中
    pub fn find_end_binary(&self, key: &[u8]) -> Option<Rc<RefCell<TrieNode<T>>>> {
        let mut current_node = self.root_node.as_ref()?.clone();

        for &byte in key.iter() {
            let next_index = byte as usize;
            let next_node = {
                let node = current_node.borrow();
                node.next[next_index].as_ref()?.clone()
            };
            current_node = next_node;
        }

        Some(current_node)
    }
}

impl<T> Trie<T> {
    /// 获取字典树中存储的条目数量
    ///
    /// # 返回值
    /// 返回 `usize` 类型的条目数量，当根节点不存在时返回 0
    pub fn num_entries(&self) -> usize {
        self.size
    }
}
impl<T> Trie<T> {
    /// 查找给定键对应的末端节点
    ///
    /// # 参数
    /// - `key`: 二进制键值，使用字节切片表示（对应 C 的 `char*` 但更安全）
    ///
    /// # 返回值
    /// 返回 Option<Rc<RefCell<TrieNode<T>>>>：
    /// - Some 包含末端节点的智能指针
    /// - None 表示键路径不完整
    pub fn find_end(&self, key: &[u8]) -> Option<Rc<RefCell<TrieNode<T>>>> {
        let mut current_node = self.root_node.as_ref().map(Rc::clone);

        for &byte in key.iter() {
            let node = match current_node.take() {
                Some(n) => n,
                None => break,
            };

            let next_index = byte as usize;
            current_node = node.borrow().next[next_index].as_ref().map(Rc::clone);
        }

        current_node
    }
}

impl<T> Trie<T> {
    /// 回滚未完成的插入操作并清理节点
    ///
    /// # 参数
    /// - `key`: 需要回滚的二进制键路径
    ///
    /// # 注意事项
    /// 遍历路径节点递减引用计数，当计数归零时自动释放节点
    /// 需要保持与原有插入操作相同的路径遍历顺序
    pub fn insert_rollback(&mut self, key: &[u8]) {
        let mut current_node = self.root_node.take();
        if current_node.is_none() {
            return;
        }
        let mut parent_link = &mut self.root_node;

        for &byte in key {
            let index = byte as usize;
            // 取出当前节点的Rc
            let node_rc = current_node.take().expect("current_node should exist here");
            let mut node = match Rc::try_unwrap(node_rc) {
                Ok(refcell) => refcell.into_inner(),
                Err(_) => panic!("Unexpected multiple references"),
            };

            // 安全递减引用计数
            node.use_count = node.use_count.saturating_sub(1);

            // 处理需要清理的节点
            if node.use_count == 0 {
                *parent_link = None;
                return;
            }

            // 获取并断开下一个节点链接
            let next_node = node.next[index].take();
            // 将处理后的节点重新存入父链接
            *parent_link = Some(Rc::new(RefCell::new(node)));

            // 更新父链接为当前节点的next数组的对应项
            parent_link = &mut Rc::get_mut(parent_link.as_mut().unwrap())
                .unwrap()
                .get_mut()
                .next[index];

            current_node = next_node;

            // 如果没有下一个节点，提前退出循环
            if current_node.is_none() {
                break;
            }
        }

        // 将剩余的current_node重新链接到父链接
        // 重建剩余节点链
        // *parent_link = Some(current_node);
    }

}
impl<T> Drop for Trie<T> {
    /// 实现字典树的非递归析构方法
    ///
    /// 使用栈式内存回收策略避免深度递归导致的栈溢出问题：
    /// 1. 将根节点压入回收栈
    /// 2. 循环弹出节点并处理其所有子节点
    /// 3. 通过 Rc::try_unwrap 确保节点独占所有权后进行析构
    /// 4. 自动管理 RefCell 内部值的释放
    fn drop(&mut self) {
        let mut stack = Vec::new();

        // 取出根节点所有权开始处理
        if let Some(root) = self.root_node.take() {
            stack.push(root);
        }

        // 栈式内存回收
        while let Some(node) = stack.pop() {
            if let Ok(node) = Rc::try_unwrap(node) {
                let mut node = node.into_inner();

                // 将子节点加入回收栈
                for child in node.next.iter_mut() {
                    if let Some(child) = child.take() {
                        stack.push(child);
                    }
                }
            }
        }
    }
}

impl<T> Trie<T> {
    /// 从字典树中移除指定二进制键对应的条目
    ///
    /// # 参数
    /// - `key`: 二进制键值，使用字节切片表示
    ///
    /// # 返回值
    /// 返回 `bool` 表示是否成功移除条目
    pub fn remove_binary(&mut self, key: &[u8]) -> bool {
        // 查找末端节点并清除数据
        let Some(end_node) = self.find_end_binary(key) else {
            return false;
        };

        // 检查末端节点是否有效
        let mut end_borrow = end_node.borrow_mut();
        if end_borrow.data.is_none() {
            return false;
        }
        end_borrow.data = None;
        self.size -= 1;
        drop(end_borrow);

        // 从根节点开始遍历路径
        let mut parent_link = &mut self.root_node;
        let Some(mut current_node) = parent_link.take() else {
            return true;
        };

        for &byte in key {
            let index = byte as usize;

            // 解包当前节点并更新引用计数
            let node = match Rc::try_unwrap(current_node) {
                Ok(n) => n.into_inner(),
                Err(rc) => {
                    // 存在共享引用时使用借用方式更新
                    let new_use_count = {
                        let mut borrowed = rc.borrow_mut();
                        borrowed.use_count = borrowed.use_count.saturating_sub(1);
                        borrowed.use_count
                    };

                    if new_use_count == 0 {
                        *parent_link = None;
                        return true;
                    }

                    *parent_link = Some(rc);
                    let parent_rc = Rc::get_mut(parent_link.as_mut().unwrap()).unwrap();
                    let parent_node = RefCell::get_mut(parent_rc);
                    parent_link = &mut parent_node.next[index];
                    current_node = parent_link.take().unwrap();
                    continue;
                }
            };

            // 更新引用计数并处理节点回收
            let mut node = node;
            node.use_count = node.use_count.saturating_sub(1);

            if node.use_count == 0 {
                *parent_link = None;
                return true;
            }

            // 重建节点链并继续遍历
            let next_node = node.next[index].take();
            *parent_link = Some(Rc::new(RefCell::new(node)));
            parent_link = &mut Rc::get_mut(parent_link.as_mut().unwrap())
                .unwrap()
                .get_mut()
                .next[index];
            current_node = match next_node {
                Some(n) => n,
                None => break,
            };
        }
        // 重建剩余节点链
        // *parent_link = Some(current_node);
        // *parent_link = Some(current_node);
        true
    }
}

impl<T: Clone> Trie<T> {
    /// 查找二进制键对应的存储值
    ///
    /// # 参数
    /// - `key`: 二进制键值，以字节切片形式传递
    ///
    /// # 返回值
    /// 返回 `TrieValue<T>`：
    /// - `Some(T)` 包含找到的值
    /// - `None` 表示键不存在或对应节点无数据
    pub fn lookup_binary(&self, key: &[u8]) -> TrieValue<T> {
        self.find_end_binary(key)
            .and_then(|node| node.borrow().data.clone())
    }
}

impl<T: Clone> Trie<T> {
    /// 查找字符串键对应的存储值
    ///
    /// # 参数
    /// - `key`: 字符串键值，自动转换为字节序列处理
    ///
    /// # 返回值
    /// 返回 `TrieValue<T>`：
    /// - `Some(T)` 包含找到的值（需实现 Clone trait）
    /// - `None` 表示键不存在或对应节点无数据
    pub fn lookup(&self, key: &str) -> TrieValue<T> {
        self.find_end(key.as_bytes())
            .and_then(|node| node.borrow().data.clone())
    }
}

impl<T> Trie<T> {
    /// 从字典树中移除指定字符串键对应的条目
    ///
    /// # 参数
    /// - `key`: 字符串键值，以字节切片形式传递（兼容C字符串格式）
    ///
    /// # 返回值
    /// 返回 `bool` 表示是否成功移除条目
    pub fn remove(&mut self, key: &[u8]) -> bool {
        // 查找末端节点并清除数据
        let Some(end_node) = self.find_end(key) else {
            return false;
        };

        if end_node.borrow().data.is_none() {
            return false;
        }
        end_node.borrow_mut().data = None;
        self.size -= 1;

        // 从根节点开始遍历路径
        let mut parent_link = &mut self.root_node;
        let Some(mut current_node) = parent_link.take() else {
            return true;
        };

        for &byte in key {
            let index = byte as usize;

            // 解包当前节点并更新引用计数
            let node = match Rc::try_unwrap(current_node) {
                Ok(n) => {
                    let mut node = n.into_inner();
                    node.use_count = node.use_count.saturating_sub(1);
                    if node.use_count == 0 {
                        *parent_link = None;
                        return true;
                    }
                    node
                }
                Err(rc) => {
                    {
                        let mut borrowed = rc.borrow_mut();
                        borrowed.use_count = borrowed.use_count.saturating_sub(1);
                        if borrowed.use_count == 0 {
                            *parent_link = None;
                            return true;
                        }
                    }
                    *parent_link = Some(rc);
                    let next_node = {
                        let rc_ref = parent_link.as_ref().unwrap();
                        let node_borrow = rc_ref.borrow();
                        node_borrow.next[index].clone()
                    };
                    current_node = match next_node {
                        Some(n) => n,
                        None => break,
                    };
                    continue;
                }
            };

            // 更新引用计数并处理节点回收
            let mut node = node;
            node.use_count = node.use_count.saturating_sub(1);

            if node.use_count == 0 {
                *parent_link = None;
                return true;
            }

            // 重建节点链并继续遍历
            let next_node = node.next[index].take();
            *parent_link = Some(Rc::new(RefCell::new(node)));
            parent_link = &mut Rc::get_mut(parent_link.as_mut().unwrap())
                .unwrap()
                .get_mut()
                .next[index];
            current_node = match next_node {
                Some(n) => n,
                None => break,
            };
        }

        true
    }
}

impl<T> Trie<T> {
    /// 向字典树插入二进制键值对
    ///
    /// # 参数
    /// - `key`: 二进制键值，使用字节切片表示
    /// - `value`: 要存储的值，`Option<T>` 类型（None 表示删除操作）
    ///
    /// # 返回值
    /// 返回 `bool` 表示是否插入成功
    pub fn insert_binary(&mut self, key: &[u8], value: Option<T>) -> bool {
        // 拒绝插入空值
        if value.is_none() {
            return false;
        }

        // 处理空键的特殊情况
        if key.is_empty() {
            if let Some(root) = self.root_node.as_ref() {
                let mut root_borrow = root.borrow_mut();
                let was_some = root_borrow.data.is_some();
                root_borrow.data = value;
                if !was_some {
                    self.size += 1;
                }
                return true;
            } else {
                self.root_node = Some(Rc::new(RefCell::new(TrieNode {
                    data: value,
                    use_count: 1,
                    next: std::array::from_fn(|_| None),
                })));
                self.size += 1;
                return true;
            }
        }

        // 检查是否已存在相同键
        // 检查是否已存在相同键
        if let Some(node) = self.find_end_binary(key) {
            let mut node_borrow = node.borrow_mut();
            let was_some = node_borrow.data.is_some();
            node_borrow.data = value;
            if !was_some {
                self.size += 1;
            }
            return true;
        }

        // 确保路径节点存在
        // 确保路径节点存在
        // 确保路径节点存在
        let mut current_link = self.root_node.clone();

        // 处理非空键时确保根节点存在
        if !key.is_empty() && self.root_node.is_none() {
            self.root_node = Some(Rc::new(RefCell::new(TrieNode {
                data: None,
                use_count: 1,
                next: std::array::from_fn(|_| None),
            })));
            current_link = self.root_node.clone();
        }

        for &byte in key {
            let index = byte as usize;

            // 获取或创建当前节点
            let current_rc = match current_link.take() {
                Some(rc) => rc,
                None => {
                    let new_node = Rc::new(RefCell::new(TrieNode {
                        data: None,
                        use_count: 1,
                        next: std::array::from_fn(|_| None),
                    }));
                    current_link = Some(new_node.clone());
                    new_node
                }
            };

            // 获取子节点或创建新的子节点
            let mut current_node = current_rc.borrow_mut();
            current_node.use_count += 1;
            // 如果子节点不存在则创建新节点
            if current_node.next[index].is_none() {
                current_node.next[index] = Some(Rc::new(RefCell::new(TrieNode {
                    data: None,
                    use_count: 1,
                    next: std::array::from_fn(|_| None),
                })));
            }

            // 移动到子节点
            current_link = current_node.next[index].clone();
        }

        // 创建末端节点并设置值
        if let Some(node_rc) = current_link {
            let mut node = node_rc.borrow_mut();
            if node.data.is_none() {
                self.size += 1;
            }
            node.data = value;
            true
        } else {
            false
        }
    }
}